#!/usr/bin/python
#coding=utf-8
'''
    4.3 实现迭代协议
    构建一个支持迭代操作的自定义对象，
    找到一个能实现迭代协议的简单方法
'''

# 以深度优先方式遍历树形节点的生成器
class Node:
    def __init__(self, value):
        self._value = value
        self._children = []

    def __repr__(self):
        return 'Node({!r})'.format(self._value)

    def add_child(self, node):
        self._children.append(node)

    def __iter__(self):
        return iter(self._children)

    def depth_first(self):
        yield self
        for c in self:
            yield from c.depth_first()

if __name__ == "__main__":
    root = Node(0)
    child1 = Node(1)
    child2 = Node(2)
    root.add_child(child1)
    root.add_child(child2)
    child1.add_child(Node(3))
    child1.add_child(Node(4))
    child2.add_child(Node(5))

    for ch in root.depth_first():
        print(ch)

## depth_first()方法：
## 它首先返回自己本身并迭代每一个子节点
## 并通过调用子节点的depth_first()方法（使用yield from 语句）返回对应元素


# Part2
## Python的迭代协议要求一个__iter__()方法 返回一个特殊的迭代器对象。
## 这个迭代器对象实现了__next__()方法，
## 并通过StopIteration异常标识迭代的完成。

## 如何使用一个关联迭代器类 重新实现depth_first() 方法
class Node2:
    def __init__(self, value):
        self._value = value
        self._children = []

    def __repr__(self):
        return 'Node({!r})'.format(self._value)

    def add_child(self, node):
        self._children.append(node)

    def __iter__(self):
        return iter(self._children)

    def depth_first(self):
        return DepthFirstIterator(self)

class DepthFirstIterator(object):
    '''
    Depth-first-traversal
    '''

    def __init__(self, start_node):
        self._node = start_node
        self._children_iter = None
        self._child_iter = None

    def __iter__(self):
        return self

    def __next__(self):
        # Return myself if just started, create an iterator for children
        if self._child_iter is None:
            self._child_iter = iter(self._node)
            return self._node

        # If processing a child, return its next item
        elif self._child_iter:
            try:
                nextchild = next(self._child_iter)
                return nextchild
            except StopIteration:
                self._child_iter = None
                return next(self)

        # Advance to the next child and start its iteration
        else:
            self._child_iter = next(self._child_iter).depth_first()
            return next(self)

